package com.lordjoe.collectiveintelligence.svm;

import libsvm.*;

import java.io.*;

import com.lordjoe.collectiveintelligence.svm.matchmaker.*;

/**
 * com.lordjoe.collectiveintelligence.svm.SVMTrainer
 *
 * @author Steve Lewis
 * @date May 24, 2009
 */
public class SVMRunner
{
    public static SVMRunner[] EMPTY_ARRAY = {};
    public static Class THIS_CLASS = SVMRunner.class;

    private svm_parameter m_Param;        // set by parse_command_line
    private svm_problem m_Problem;        // set by read_problem
    private svm_model m_Model;
    private String model_file_name;        // set by parse_command_line
    private String error_msg;
    private int cross_validation;
    private int nr_fold;


    public SVMRunner(ISVMDataSource ds)
    {
        m_Param = SVMUtilities.buildDefaultParameters();
        m_Problem = SVMUtilities.buildSVMProblem(ds);

    }

    public svm_parameter getParam()
    {
        return m_Param;
    }

    public void setParam(svm_parameter pParam)
    {
        m_Param = pParam;
    }

    public svm_problem getProblem()
    {
        return m_Problem;
    }

    public void setProblem(svm_problem pProblem)
    {
        m_Problem = pProblem;
    }

    public svm_model getModel()
    {
        return m_Model;
    }

    public void setModel(svm_model pModel)
    {
        m_Model = pModel;
    }

    public String getModel_file_name()
    {
        return model_file_name;
    }

    public void setModel_file_name(String pModel_file_name)
    {
        model_file_name = pModel_file_name;
    }

    public void predict(ISVMDataSource ds)
    {
        svm_model model = getModel();
        int svm_type = svm.svm_get_svm_type(model);
        int nr_class = svm.svm_get_nr_class(model);
        ISVMDataItem[] items = ds.getDataItems();
        int numberBad = 0;
        for (int i = 0; i < items.length; i++) {
            double[] prob_estimates = new double[nr_class];
            ISVMDataItem item = items[i];
            svm_node[] nodes = item.getNodes();
            double predict = svm.svm_predict_probability(model, nodes, prob_estimates);
            //  double prob  = svm.svm_predict(model,nodes);
            // if (prob != -1)
            item.setPredictedClassification(predict);
            System.out.print(predict);
            System.out.print(" ");
            double v = item.getClassification();
            if(v != predict)    {
                 System.out.println(" BAD");
                numberBad++;
            }
            else
                System.out.println();


        }
        System.out.println("Number Bad " + numberBad + " Number " + items.length + 
          " Fraction bad " + ((double)numberBad) / items.length);
    }

    private static void matchmakerProblem() throws IOException
    {
        String fileBase = MatchMaker.DATA_PATH + "MatchMaker";
        MatchMaker mm = MatchMaker.buildMatchmaker(MatchMaker.DATA_PATH + "MatchMaker.tab");
        MatchMakerDataSource dataSource = new MatchMakerDataSource(DataSetType.Raw, mm);
        ISVMDataSource[] sets = SVMUtilities.buildTrainingAndProductionSets(dataSource,
                dataSource.getDataItems().length / 2);
        ISVMDataSource training = sets[0];
        ISVMDataSource production = sets[1];
        
        SVMUtilities.saveDataSource(training,fileBase + ".training");
        SVMUtilities.saveDataSource(production,fileBase + ".model");

        SVMRunner trainer = new SVMRunner(training);
        String modelFileName = fileBase + ".computedModel";
        trainer.setModel_file_name(modelFileName);
        File modelFile = new File(modelFileName);
        svm_model svm_model = null;
        if (false && modelFile.exists()) {
            svm_model = svm.svm_load_model(modelFileName);
        }
        else {

            System.out.println("Ready to Train");
            svm_problem svm_problem = trainer.getProblem();
            svm_parameter param = trainer.getParam();
            param.C = 128;
            svm_model = svm.svm_train(svm_problem, param);
            svm.svm_save_model(modelFileName, svm_model);
        }
        trainer.setModel(svm_model);
        trainer.predict(production);

        PotentialPair[] pairs = mm.getPotentialPairs();

        MatchMaker.showAgePredictionEffect(pairs);
         MatchMaker.showDistanceAndInterestsPredictionEffect(pairs);
        
    }


    private static void trainOnDataset(String pBaseFile)
            throws IOException
    {
        ISVMDataSource[] sets = SVMUtilities.readTrainingAndModelSets(pBaseFile);
        ISVMDataSource training = sets[0];
        ISVMDataSource production = sets[1];
        SVMRunner trainer = new SVMRunner(training);
        String modelFileName = pBaseFile + ".computedModel";
        trainer.setModel_file_name(modelFileName);
        File modelFile = new File(modelFileName);
        svm_model svm_model = null;
        if (modelFile.exists()) {
            svm_model = svm.svm_load_model(modelFileName);
        }
        else {
            System.out.println("Ready to Train");
            svm_model = svm.svm_train(trainer.getProblem(), trainer.getParam());
            svm.svm_save_model(modelFileName, svm_model);
        }
        trainer.setModel(svm_model);

        trainer.predict(production);
    }

    public static void main(String[] args) throws IOException
    {
        if (args.length == 0) {
            String[] baseFile = {"data/svm/a3a"};
            args = baseFile;
        }
        if (args[0].equals("matchmaker")) {
            matchmakerProblem();
            return;
        }
        for (int i = 0; i < args.length; i++) {
            String arg = args[i];
            System.out.println(arg);
            trainOnDataset(arg);

        }
    }


}
